{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据处理\n",
    "\n",
    "开始实验前，需确保本地已经安装Python环境及MindSpore。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0.0a0\n"
     ]
    }
   ],
   "source": [
    "import mindspore\n",
    "print(mindspore.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating data folder...\n",
      "Downloading data from https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/dataset_fcn8s.tar (537.2 MB)\n",
      "\n",
      "file_sizes: 100%|█████████████████████████████| 563M/563M [00:02<00:00, 201MB/s]\n",
      "Extracting tar file...\n",
      "Successfully downloaded / unzipped to ./dataset\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'./dataset'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from download import download\n",
    "\n",
    "url = \"https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/dataset_fcn8s.tar\"\n",
    "\n",
    "download(url, \"./dataset\", kind=\"tar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据预处理\n",
    "\n",
    "由于PASCAL VOC 2012数据集中图像的分辨率大多不一致，无法放在一个tensor中，故输入前需做标准化处理。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据加载\n",
    "\n",
    "将PASCAL VOC 2012数据集与SDB数据集进行混合。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "import mindspore.dataset as ds\n",
    "\n",
    "class SegDataset:\n",
    "    def __init__(self,\n",
    "                 image_mean,\n",
    "                 image_std,\n",
    "                 data_file='',\n",
    "                 batch_size=32,\n",
    "                 crop_size=512,\n",
    "                 max_scale=2.0,\n",
    "                 min_scale=0.5,\n",
    "                 ignore_label=255,\n",
    "                 num_classes=21,\n",
    "                 num_readers=2,\n",
    "                 num_parallel_calls=4):\n",
    "\n",
    "        self.data_file = data_file\n",
    "        self.batch_size = batch_size\n",
    "        self.crop_size = crop_size\n",
    "        self.image_mean = np.array(image_mean, dtype=np.float32)\n",
    "        self.image_std = np.array(image_std, dtype=np.float32)\n",
    "        self.max_scale = max_scale\n",
    "        self.min_scale = min_scale\n",
    "        self.ignore_label = ignore_label\n",
    "        self.num_classes = num_classes\n",
    "        self.num_readers = num_readers\n",
    "        self.num_parallel_calls = num_parallel_calls\n",
    "        max_scale > min_scale\n",
    "\n",
    "    def preprocess_dataset(self, image, label):\n",
    "        image_out = cv2.imdecode(np.frombuffer(image, dtype=np.uint8), cv2.IMREAD_COLOR)\n",
    "        label_out = cv2.imdecode(np.frombuffer(label, dtype=np.uint8), cv2.IMREAD_GRAYSCALE)\n",
    "        sc = np.random.uniform(self.min_scale, self.max_scale)\n",
    "        new_h, new_w = int(sc * image_out.shape[0]), int(sc * image_out.shape[1])\n",
    "        image_out = cv2.resize(image_out, (new_w, new_h), interpolation=cv2.INTER_CUBIC)\n",
    "        label_out = cv2.resize(label_out, (new_w, new_h), interpolation=cv2.INTER_NEAREST)\n",
    "\n",
    "        image_out = (image_out - self.image_mean) / self.image_std\n",
    "        out_h, out_w = max(new_h, self.crop_size), max(new_w, self.crop_size)\n",
    "        pad_h, pad_w = out_h - new_h, out_w - new_w\n",
    "        if pad_h > 0 or pad_w > 0:\n",
    "            image_out = cv2.copyMakeBorder(image_out, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0)\n",
    "            label_out = cv2.copyMakeBorder(label_out, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=self.ignore_label)\n",
    "        offset_h = np.random.randint(0, out_h - self.crop_size + 1)\n",
    "        offset_w = np.random.randint(0, out_w - self.crop_size + 1)\n",
    "        image_out = image_out[offset_h: offset_h + self.crop_size, offset_w: offset_w + self.crop_size, :]\n",
    "        label_out = label_out[offset_h: offset_h + self.crop_size, offset_w: offset_w+self.crop_size]\n",
    "        if np.random.uniform(0.0, 1.0) > 0.5:\n",
    "            image_out = image_out[:, ::-1, :]\n",
    "            label_out = label_out[:, ::-1]\n",
    "        image_out = image_out.transpose((2, 0, 1))\n",
    "        image_out = image_out.copy()\n",
    "        label_out = label_out.copy()\n",
    "        label_out = label_out.astype(\"int32\")\n",
    "        return image_out, label_out\n",
    "\n",
    "    def get_dataset(self):\n",
    "        ds.config.set_numa_enable(True)\n",
    "        dataset = ds.MindDataset(self.data_file, columns_list=[\"data\", \"label\"],\n",
    "                                 shuffle=True, num_parallel_workers=self.num_readers)\n",
    "        transforms_list = self.preprocess_dataset\n",
    "        dataset = dataset.map(operations=transforms_list, input_columns=[\"data\", \"label\"],\n",
    "                              output_columns=[\"data\", \"label\"],\n",
    "                              num_parallel_workers=self.num_parallel_calls)\n",
    "        dataset = dataset.shuffle(buffer_size=self.batch_size * 10)\n",
    "        dataset = dataset.batch(self.batch_size, drop_remainder=True)\n",
    "        return dataset\n",
    "\n",
    "\n",
    "# 定义创建数据集的参数\n",
    "IMAGE_MEAN = [103.53, 116.28, 123.675]\n",
    "IMAGE_STD = [57.375, 57.120, 58.395]\n",
    "DATA_FILE = \"dataset/dataset_fcn8s/mindname.mindrecord\"\n",
    "\n",
    "# 定义模型训练参数\n",
    "train_batch_size = 4\n",
    "crop_size = 512\n",
    "min_scale = 0.5\n",
    "max_scale = 2.0\n",
    "ignore_label = 255\n",
    "num_classes = 21\n",
    "\n",
    "# 实例化Dataset\n",
    "dataset = SegDataset(image_mean=IMAGE_MEAN,\n",
    "                     image_std=IMAGE_STD,\n",
    "                     data_file=DATA_FILE,\n",
    "                     batch_size=train_batch_size,\n",
    "                     crop_size=crop_size,\n",
    "                     max_scale=max_scale,\n",
    "                     min_scale=min_scale,\n",
    "                     ignore_label=ignore_label,\n",
    "                     num_classes=num_classes,\n",
    "                     num_readers=2,\n",
    "                     num_parallel_calls=4)\n",
    "\n",
    "dataset = dataset.get_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 训练集可视化\n",
    "\n",
    "运行以下代码观察载入的数据集图片（数据处理过程中已做归一化处理）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(16, 8))\n",
    "\n",
    "# 对训练集中的数据进行展示\n",
    "for i in range(1, 9):\n",
    "    plt.subplot(2, 4, i)\n",
    "    show_data = next(dataset.create_dict_iterator())\n",
    "    show_images = show_data[\"data\"].asnumpy()\n",
    "    show_images = np.clip(show_images, 0, 1)\n",
    "# 将图片转换HWC格式后进行展示\n",
    "    plt.imshow(show_images[0].transpose(1, 2, 0))\n",
    "    plt.axis(\"off\")\n",
    "    plt.subplots_adjust(wspace=0.05, hspace=0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 网络构建\n",
    "\n",
    "使用以下代码构建FCN-8s网络。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore.nn as nn\n",
    "\n",
    "class FCN8s(nn.Cell):\n",
    "    def __init__(self, n_class):\n",
    "        super().__init__()\n",
    "        self.n_class = n_class\n",
    "        self.conv1 = nn.SequentialCell(\n",
    "            nn.Conv2d(in_channels=3, out_channels=64,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=64, out_channels=64,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.conv2 = nn.SequentialCell(\n",
    "            nn.Conv2d(in_channels=64, out_channels=128,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=128, out_channels=128,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.conv3 = nn.SequentialCell(\n",
    "            nn.Conv2d(in_channels=128, out_channels=256,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=256, out_channels=256,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=256, out_channels=256,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.conv4 = nn.SequentialCell(\n",
    "            nn.Conv2d(in_channels=256, out_channels=512,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=512, out_channels=512,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=512, out_channels=512,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.conv5 = nn.SequentialCell(\n",
    "            nn.Conv2d(in_channels=512, out_channels=512,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=512, out_channels=512,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels=512, out_channels=512,\n",
    "                      kernel_size=3, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.conv6 = nn.SequentialCell(\n",
    "            nn.Conv2d(in_channels=512, out_channels=4096,\n",
    "                      kernel_size=7, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(4096),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.conv7 = nn.SequentialCell(\n",
    "            nn.Conv2d(in_channels=4096, out_channels=4096,\n",
    "                      kernel_size=1, weight_init='xavier_uniform'),\n",
    "            nn.BatchNorm2d(4096),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.score_fr = nn.Conv2d(in_channels=4096, out_channels=self.n_class,\n",
    "                                  kernel_size=1, weight_init='xavier_uniform')\n",
    "        self.upscore2 = nn.Conv2dTranspose(in_channels=self.n_class, out_channels=self.n_class,\n",
    "                                           kernel_size=4, stride=2, weight_init='xavier_uniform')\n",
    "        self.score_pool4 = nn.Conv2d(in_channels=512, out_channels=self.n_class,\n",
    "                                     kernel_size=1, weight_init='xavier_uniform')\n",
    "        self.upscore_pool4 = nn.Conv2dTranspose(in_channels=self.n_class, out_channels=self.n_class,\n",
    "                                                kernel_size=4, stride=2, weight_init='xavier_uniform')\n",
    "        self.score_pool3 = nn.Conv2d(in_channels=256, out_channels=self.n_class,\n",
    "                                     kernel_size=1, weight_init='xavier_uniform')\n",
    "        self.upscore8 = nn.Conv2dTranspose(in_channels=self.n_class, out_channels=self.n_class,\n",
    "                                           kernel_size=16, stride=8, weight_init='xavier_uniform')\n",
    "\n",
    "    def construct(self, x):\n",
    "        x1 = self.conv1(x)\n",
    "        p1 = self.pool1(x1)\n",
    "        x2 = self.conv2(p1)\n",
    "        p2 = self.pool2(x2)\n",
    "        x3 = self.conv3(p2)\n",
    "        p3 = self.pool3(x3)\n",
    "        x4 = self.conv4(p3)\n",
    "        p4 = self.pool4(x4)\n",
    "        x5 = self.conv5(p4)\n",
    "        p5 = self.pool5(x5)\n",
    "        x6 = self.conv6(p5)\n",
    "        x7 = self.conv7(x6)\n",
    "        sf = self.score_fr(x7)\n",
    "        u2 = self.upscore2(sf)\n",
    "        s4 = self.score_pool4(p4)\n",
    "        f4 = s4 + u2\n",
    "        u4 = self.upscore_pool4(f4)\n",
    "        s3 = self.score_pool3(p3)\n",
    "        f3 = s3 + u4\n",
    "        out = self.upscore8(f3)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练准备\n",
    "\n",
    "### 导入VGG-16部分预训练权重\n",
    "\n",
    "FCN使用VGG-16作为骨干网络，用于实现图像编码。使用下面代码导入VGG-16预训练模型的部分预训练权重。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Replace is False and data exists, so doing nothing. Use replace=True to re-download the data.\n"
     ]
    }
   ],
   "source": [
    "from download import download\n",
    "from mindspore import load_checkpoint, load_param_into_net\n",
    "\n",
    "url = \"https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/fcn8s_vgg16_pretrain.ckpt\"\n",
    "download(url, \"fcn8s_vgg16_pretrain.ckpt\")\n",
    "def load_vgg16():\n",
    "    ckpt_vgg16 = \"fcn8s_vgg16_pretrain.ckpt\"\n",
    "    param_vgg = load_checkpoint(ckpt_vgg16)\n",
    "    load_param_into_net(net, param_vgg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 自定义评价指标 Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import mindspore as ms\n",
    "import mindspore.nn as nn\n",
    "import mindspore.train as train\n",
    "\n",
    "class PixelAccuracy(train.Metric):\n",
    "    def __init__(self, num_class=21):\n",
    "        super(PixelAccuracy, self).__init__()\n",
    "        self.num_class = num_class\n",
    "\n",
    "    def _generate_matrix(self, gt_image, pre_image):\n",
    "        mask = (gt_image >= 0) & (gt_image < self.num_class)\n",
    "        label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]\n",
    "        count = np.bincount(label, minlength=self.num_class**2)\n",
    "        confusion_matrix = count.reshape(self.num_class, self.num_class)\n",
    "        return confusion_matrix\n",
    "\n",
    "    def clear(self):\n",
    "        self.confusion_matrix = np.zeros((self.num_class,) * 2)\n",
    "\n",
    "    def update(self, *inputs):\n",
    "        y_pred = inputs[0].asnumpy().argmax(axis=1)\n",
    "        y = inputs[1].asnumpy().reshape(4, 512, 512)\n",
    "        self.confusion_matrix += self._generate_matrix(y, y_pred)\n",
    "\n",
    "    def eval(self):\n",
    "        pixel_accuracy = np.diag(self.confusion_matrix).sum() / self.confusion_matrix.sum()\n",
    "        return pixel_accuracy\n",
    "\n",
    "\n",
    "class PixelAccuracyClass(train.Metric):\n",
    "    def __init__(self, num_class=21):\n",
    "        super(PixelAccuracyClass, self).__init__()\n",
    "        self.num_class = num_class\n",
    "\n",
    "    def _generate_matrix(self, gt_image, pre_image):\n",
    "        mask = (gt_image >= 0) & (gt_image < self.num_class)\n",
    "        label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]\n",
    "        count = np.bincount(label, minlength=self.num_class**2)\n",
    "        confusion_matrix = count.reshape(self.num_class, self.num_class)\n",
    "        return confusion_matrix\n",
    "\n",
    "    def update(self, *inputs):\n",
    "        y_pred = inputs[0].asnumpy().argmax(axis=1)\n",
    "        y = inputs[1].asnumpy().reshape(4, 512, 512)\n",
    "        self.confusion_matrix += self._generate_matrix(y, y_pred)\n",
    "\n",
    "    def clear(self):\n",
    "        self.confusion_matrix = np.zeros((self.num_class,) * 2)\n",
    "\n",
    "    def eval(self):\n",
    "        mean_pixel_accuracy = np.diag(self.confusion_matrix) / self.confusion_matrix.sum(axis=1)\n",
    "        mean_pixel_accuracy = np.nanmean(mean_pixel_accuracy)\n",
    "        return mean_pixel_accuracy\n",
    "\n",
    "\n",
    "class MeanIntersectionOverUnion(train.Metric):\n",
    "    def __init__(self, num_class=21):\n",
    "        super(MeanIntersectionOverUnion, self).__init__()\n",
    "        self.num_class = num_class\n",
    "\n",
    "    def _generate_matrix(self, gt_image, pre_image):\n",
    "        mask = (gt_image >= 0) & (gt_image < self.num_class)\n",
    "        label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]\n",
    "        count = np.bincount(label, minlength=self.num_class**2)\n",
    "        confusion_matrix = count.reshape(self.num_class, self.num_class)\n",
    "        return confusion_matrix\n",
    "\n",
    "    def update(self, *inputs):\n",
    "        y_pred = inputs[0].asnumpy().argmax(axis=1)\n",
    "        y = inputs[1].asnumpy().reshape(4, 512, 512)\n",
    "        self.confusion_matrix += self._generate_matrix(y, y_pred)\n",
    "\n",
    "    def clear(self):\n",
    "        self.confusion_matrix = np.zeros((self.num_class,) * 2)\n",
    "\n",
    "    def eval(self):\n",
    "        mean_iou = np.diag(self.confusion_matrix) / (\n",
    "            np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -\n",
    "            np.diag(self.confusion_matrix))\n",
    "        mean_iou = np.nanmean(mean_iou)\n",
    "        return mean_iou\n",
    "\n",
    "\n",
    "class FrequencyWeightedIntersectionOverUnion(train.Metric):\n",
    "    def __init__(self, num_class=21):\n",
    "        super(FrequencyWeightedIntersectionOverUnion, self).__init__()\n",
    "        self.num_class = num_class\n",
    "\n",
    "    def _generate_matrix(self, gt_image, pre_image):\n",
    "        mask = (gt_image >= 0) & (gt_image < self.num_class)\n",
    "        label = self.num_class * gt_image[mask].astype('int') + pre_image[mask]\n",
    "        count = np.bincount(label, minlength=self.num_class**2)\n",
    "        confusion_matrix = count.reshape(self.num_class, self.num_class)\n",
    "        return confusion_matrix\n",
    "\n",
    "    def update(self, *inputs):\n",
    "        y_pred = inputs[0].asnumpy().argmax(axis=1)\n",
    "        y = inputs[1].asnumpy().reshape(4, 512, 512)\n",
    "        self.confusion_matrix += self._generate_matrix(y, y_pred)\n",
    "\n",
    "    def clear(self):\n",
    "        self.confusion_matrix = np.zeros((self.num_class,) * 2)\n",
    "\n",
    "    def eval(self):\n",
    "        freq = np.sum(self.confusion_matrix, axis=1) / np.sum(self.confusion_matrix)\n",
    "        iu = np.diag(self.confusion_matrix) / (\n",
    "            np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -\n",
    "            np.diag(self.confusion_matrix))\n",
    "\n",
    "        frequency_weighted_iou = (freq[freq > 0] * iu[freq > 0]).sum()\n",
    "        return frequency_weighted_iou"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型训练\n",
    "\n",
    "导入VGG-16预训练参数后，实例化损失函数、优化器，使用Model接口编译网络，训练FCN-8s网络。\n",
    "ps：如果可以连续训练则运行1和4，一次完成全部训练；如果会间断则运行2和3，从间断前保存的最新检查点继续训练；"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#1\n",
    "import mindspore\n",
    "from mindspore import Tensor\n",
    "import mindspore.nn as nn\n",
    "from mindspore.train import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor, Model\n",
    "\n",
    "device_target = \"GPU\"\n",
    "mindspore.set_context(mode=mindspore.PYNATIVE_MODE, device_target=device_target)\n",
    "\n",
    "train_batch_size = 4\n",
    "num_classes = 21\n",
    "# 初始化模型结构\n",
    "net = FCN8s(n_class=21)\n",
    "# 导入vgg16预训练参数\n",
    "load_vgg16()\n",
    "# 计算学习率\n",
    "min_lr = 0.0005\n",
    "base_lr = 0.05\n",
    "train_epochs = 100\n",
    "iters_per_epoch = dataset.get_dataset_size()\n",
    "total_step = iters_per_epoch * train_epochs\n",
    "\n",
    "lr_scheduler = mindspore.nn.cosine_decay_lr(min_lr,\n",
    "                                            base_lr,\n",
    "                                            total_step,\n",
    "                                            iters_per_epoch,\n",
    "                                            decay_epoch=2)\n",
    "lr = Tensor(lr_scheduler[-1])\n",
    "\n",
    "# 定义损失函数\n",
    "loss = nn.CrossEntropyLoss(ignore_index=255)\n",
    "# 定义优化器\n",
    "optimizer = nn.Momentum(params=net.trainable_params(), learning_rate=lr, momentum=0.9, weight_decay=0.0001)\n",
    "# 定义loss_scale\n",
    "scale_factor = 4\n",
    "scale_window = 3000\n",
    "loss_scale_manager = ms.amp.DynamicLossScaleManager(scale_factor, scale_window)\n",
    "# 初始化模型\n",
    "if device_target == \"Ascend\":\n",
    "    model = Model(net, loss_fn=loss, optimizer=optimizer, loss_scale_manager=loss_scale_manager, metrics={\"pixel accuracy\": PixelAccuracy(), \"mean pixel accuracy\": PixelAccuracyClass(), \"mean IoU\": MeanIntersectionOverUnion(), \"frequency weighted IoU\": FrequencyWeightedIntersectionOverUnion()})\n",
    "else:\n",
    "    model = Model(net, loss_fn=loss, optimizer=optimizer, metrics={\"pixel accuracy\": PixelAccuracy(), \"mean pixel accuracy\": PixelAccuracyClass(), \"mean IoU\": MeanIntersectionOverUnion(), \"frequency weighted IoU\": FrequencyWeightedIntersectionOverUnion()})\n",
    "\n",
    "# 设置ckpt文件保存的参数\n",
    "time_callback = TimeMonitor(data_size=iters_per_epoch)\n",
    "loss_callback = LossMonitor()\n",
    "callbacks = [time_callback, loss_callback]\n",
    "save_steps = 330\n",
    "config_ckpt = CheckpointConfig(save_checkpoint_steps=1,\n",
    "                               keep_checkpoint_max=None)\n",
    "ckpt_callback = ModelCheckpoint(prefix=\"FCN8s\",\n",
    "                                directory=\"./ckpt\",\n",
    "                                config=config_ckpt)\n",
    "callbacks.append(ckpt_callback)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#2\n",
    "model.train(train_epochs, dataset, callbacks=callbacks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#3\n",
    "#根据检查点继续训练\n",
    "import mindspore\n",
    "from mindspore import Tensor\n",
    "import mindspore.nn as nn\n",
    "from mindspore.train import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor, Model\n",
    "from mindspore.train.serialization import load_checkpoint, load_param_into_net\n",
    "device_target = \"GPU\"\n",
    "mindspore.set_context(mode=mindspore.PYNATIVE_MODE, device_target=device_target)\n",
    "\n",
    "train_batch_size = 4\n",
    "num_classes = 21\n",
    "# 初始化模型结构\n",
    "net = FCN8s(n_class=21)\n",
    "\n",
    "# 计算学习率\n",
    "min_lr = 0.0005\n",
    "base_lr = 0.05\n",
    "train_epochs = 98\n",
    "iters_per_epoch = dataset.get_dataset_size()\n",
    "total_step = iters_per_epoch * train_epochs\n",
    "lr_scheduler = mindspore.nn.cosine_decay_lr(min_lr,\n",
    "                                            base_lr,\n",
    "                                            total_step,\n",
    "                                            iters_per_epoch,\n",
    "                                            decay_epoch=2)\n",
    "lr = Tensor(lr_scheduler[-1])\n",
    "\n",
    "# 定义损失函数\n",
    "loss = nn.CrossEntropyLoss(ignore_index=255)\n",
    "# 定义优化器\n",
    "optimizer = nn.Momentum(params=net.trainable_params(), learning_rate=lr, momentum=0.9, weight_decay=0.0001)\n",
    "# 定义loss_scale\n",
    "scale_factor = 4\n",
    "scale_window = 3000\n",
    "loss_scale_manager = ms.amp.DynamicLossScaleManager(scale_factor, scale_window)\n",
    "# 初始化模型\n",
    "if device_target == \"Ascend\":\n",
    "    model = Model(net, loss_fn=loss, optimizer=optimizer, loss_scale_manager=loss_scale_manager, metrics={\"pixel accuracy\": PixelAccuracy(), \"mean pixel accuracy\": PixelAccuracyClass(), \"mean IoU\": MeanIntersectionOverUnion(), \"frequency weighted IoU\": FrequencyWeightedIntersectionOverUnion()})\n",
    "else:\n",
    "    model = Model(net, loss_fn=loss, optimizer=optimizer, metrics={\"pixel accuracy\": PixelAccuracy(), \"mean pixel accuracy\": PixelAccuracyClass(), \"mean IoU\": MeanIntersectionOverUnion(), \"frequency weighted IoU\": FrequencyWeightedIntersectionOverUnion()})\n",
    "\n",
    "ckpt_path = \"ckpt/FCN8s_1-5_4.ckpt\"#（修改为最新检查点）\n",
    "# 加载检查点文件\n",
    "param_dict = load_checkpoint(ckpt_path)\n",
    "# 将参数加载到网络中\n",
    "load_param_into_net(net, param_dict)\n",
    "\n",
    "# 设置ckpt文件保存的参数\n",
    "time_callback = TimeMonitor(data_size=iters_per_epoch)\n",
    "loss_callback = LossMonitor()\n",
    "callbacks = [time_callback, loss_callback]\n",
    "save_steps = 330\n",
    "config_ckpt = CheckpointConfig(save_checkpoint_steps=1,\n",
    "                               keep_checkpoint_max=None)\n",
    "ckpt_callback = ModelCheckpoint(prefix=\"FCN8s\",\n",
    "                                directory=\"./ckpt\",\n",
    "                                config=config_ckpt)\n",
    "callbacks.append(ckpt_callback)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#4\n",
    "#从检查点继续训练\n",
    "start_epoch = 3\n",
    "\n",
    "model.train(train_epochs,\n",
    "            dataset,\n",
    "            callbacks=callbacks,\n",
    "            dataset_sink_mode=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型评估\n",
    "需要修改对应检查点文件路径及名称"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#a\n",
    "# 下载已训练好的权重文件\n",
    "from download import download\n",
    "url = \"https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/FCN8s.ckpt\"\n",
    "download(url, \"FCN8s.ckpt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_MEAN = [103.53, 116.28, 123.675]\n",
    "IMAGE_STD = [57.375, 57.120, 58.395]\n",
    "DATA_FILE = \"dataset/dataset_fcn8s/mindname.mindrecord\"#这里有可能需要修改文件路径，视下载情况\n",
    "from download import download\n",
    "net = FCN8s(n_class=num_classes)\n",
    "\n",
    "ckpt_file = \"ckpt/FCN8s-1_1143.ckpt\"#这里需要更改路径为ckpt文件夹中最新的训练检查点文件/训练下载文件\n",
    "param_dict = load_checkpoint(ckpt_file)\n",
    "load_param_into_net(net, param_dict)\n",
    "\n",
    "if device_target == \"Ascend\":\n",
    "    model = Model(net, loss_fn=loss, optimizer=optimizer, loss_scale_manager=loss_scale_manager, metrics={\"pixel accuracy\": PixelAccuracy(), \"mean pixel accuracy\": PixelAccuracyClass(), \"mean IoU\": MeanIntersectionOverUnion(), \"frequency weighted IoU\": FrequencyWeightedIntersectionOverUnion()})\n",
    "else:\n",
    "    model = Model(net, loss_fn=loss, optimizer=optimizer, metrics={\"pixel accuracy\": PixelAccuracy(), \"mean pixel accuracy\": PixelAccuracyClass(), \"mean IoU\": MeanIntersectionOverUnion(), \"frequency weighted IoU\": FrequencyWeightedIntersectionOverUnion()})\n",
    "\n",
    "# 实例化Dataset\n",
    "dataset = SegDataset(image_mean=IMAGE_MEAN,\n",
    "                     image_std=IMAGE_STD,\n",
    "                     data_file=DATA_FILE,\n",
    "                     batch_size=train_batch_size,\n",
    "                     crop_size=crop_size,\n",
    "                     max_scale=max_scale,\n",
    "                     min_scale=min_scale,\n",
    "                     ignore_label=ignore_label,\n",
    "                     num_classes=num_classes,\n",
    "                     num_readers=2,\n",
    "                     num_parallel_calls=4)\n",
    "dataset_eval = dataset.get_dataset()\n",
    "model.eval(dataset_eval)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型推理\n",
    "\n",
    "使用训练的网络对模型推理结果进行展示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "net = FCN8s(n_class=num_classes)\n",
    "# 设置超参\n",
    "ckpt_file = \"ckpt/FCN8s-1_1143.ckpt\"#同样需要更改路径为ckpt文件夹中最新的训练检查点文件\n",
    "param_dict = load_checkpoint(ckpt_file)\n",
    "load_param_into_net(net, param_dict)\n",
    "eval_batch_size = 4\n",
    "img_lst = []\n",
    "mask_lst = []\n",
    "res_lst = []\n",
    "# 推理效果展示(上方为输入图片，下方为推理效果图片)\n",
    "plt.figure(figsize=(8, 5))\n",
    "show_data = next(dataset_eval.create_dict_iterator())\n",
    "show_images = show_data[\"data\"].asnumpy()\n",
    "mask_images = show_data[\"label\"].reshape([4, 512, 512])\n",
    "show_images = np.clip(show_images, 0, 1)\n",
    "for i in range(eval_batch_size):\n",
    "    img_lst.append(show_images[i])\n",
    "    mask_lst.append(mask_images[i])\n",
    "res = net(show_data[\"data\"]).asnumpy().argmax(axis=1)\n",
    "for i in range(eval_batch_size):\n",
    "    plt.subplot(2, 4, i + 1)\n",
    "    plt.imshow(img_lst[i].transpose(1, 2, 0))\n",
    "    plt.axis(\"off\")\n",
    "    plt.subplots_adjust(wspace=0.05, hspace=0.02)\n",
    "    plt.subplot(2, 4, i + 5)\n",
    "    plt.imshow(res[i])\n",
    "    plt.axis(\"off\")\n",
    "    plt.subplots_adjust(wspace=0.05, hspace=0.02)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型预测\n",
    "\n",
    "使用训练的网络对模型预测输入图片的结果进行展示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "import mindspore.numpy as mnp\n",
    "import mindspore.nn as nn\n",
    "import mindspore.ops.operations as P\n",
    "import mindspore.dataset.vision.py_transforms as transforms\n",
    "from mindspore.dataset.transforms.py_transforms import Compose\n",
    "\n",
    "# 构建FCN8s模型\n",
    "fcn_model = FCN8s(n_class=num_classes).to_float(mnp.float32)\n",
    "\n",
    "# 加载已经训练好的模型参数\n",
    "param_dict = load_checkpoint(\"ckpt/FCN8s-1_1143.ckpt\")#需要更改路径为ckpt文件夹中最新的训练检查点文件\n",
    "\n",
    "# 将参数加载到模型中\n",
    "load_param_into_net(fcn_model, param_dict)\n",
    "\n",
    "# 读入一张测试图片并进行预处理\n",
    "img = Image.open('test4.jpg')  #更改路径和名称对应文件夹中想要输入进行预测的图片\n",
    "transform = Compose([transforms.Resize((224, 224)),\n",
    "                                transforms.ToTensor(),\n",
    "                                transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                                     std=[0.229, 0.224, 0.225])])\n",
    "img = transform(img)\n",
    "img=np.array(img)\n",
    "# img=img[np.newaxis,:]\n",
    "img=mindspore.Tensor(img)\n",
    "# 将输入图片传入模型进行预测\n",
    "output = fcn_model(img)\n",
    "# 对预测结果进行后处理得到最终的分割结果\n",
    "pred_mask = output.asnumpy().argmax(axis=1)[0]\n",
    "mask_img = Image.fromarray(pred_mask.astype(np.uint8))\n",
    "mask_img.save(\"result.jpg\")\n",
    "#展示结果\n",
    "img=Image.open(\"result.jpg\")\n",
    "plt.imshow(img)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore",
   "language": "python",
   "name": "mindspore"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "5991bb6ddd18f5f28b35941317add6d392f35bceeeac2e7b1aa1f8654ae6ac72"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
